[MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE by a-sidorova · Pull Request #4620 · ROCm/aiter

a-sidorova · 2026-08-07T06:15:40Z

Motivation

Enable the tanh-approximation GELU activation (gelu_tanh, 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))))) in the AITER Composable Kernel XDL 2-stage MoE path. The MoE gridwise epilogue supports only silu/gelu/swiglu; this ports CK's gelu_tanh_and_mul (ROCm/rocm-libraries#9396) into AITER so models whose MoE experts use the GELU tanh approximation (e.g. Gemma-family MoE) can run on this path.

JIRA ID : ROCM-27619

Technical Details

Wires a new GeluTanh activation end-to-end into the CK 2-stage MoE codegen, following the existing swiglu_oai port (78e45124):

csrc/include/aiter_enum.h: add ActivationType::GeluTanh = 4.
csrc/include/rocm_ops.hpp: expose GeluTanh on the pybind ActivationType enum.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu: map ActivationType::GeluTanh -> 4 (CK gelu_tanh_and_mul) in map_activation_to_ck_stage1.
gemm_moe_ck2stages_common.py: add "gelutanh": 4 to ACT_OP_MAP (drives CK ActOP codegen and kernel-instance naming).
gen_instances.py: add gelutanh to the -act choices and a targeted plain-f8 AOT prebuild loop (mirrors the swiglu loop; not required for correctness since runtime JIT generates on demand).

CK submodule bump

Bumps 3rdparty/composable_kernel to the merged feat(ck): Added Gelu with Tanh approx to XDL 2-stage MoE epilogue rocm-libraries#9396, which provides the CK-side gelu with tanh approx epilogue. This PR must be built / merged together with that CK commit.

Test/reference support so the change can be validated:

aiter/ops/quant.py (get_torch_act): add the GeluTanh reference using the tanh-approx GELU (F.gelu(x, approximate="tanh")) to match CK's FastGelu epilogue.
aiter/utility/dtypes.py (str2ActivationType): case-insensitive enum lookup so mixed-case members like GeluTanh resolve from CLI (-a gelutanh).

The activation is applied in fp32 in the stage-1 epilogue; GEMM compute and quantization are untouched. CK-side support comes from ROCm/rocm-libraries#9396 (gelu_tanh_and_mul = 4).

Test Plan

Run the CK 2-stage MoE test with the new activation on the plain-f8 (a8w8) path:

python3 op_tests/test_moe_2stage.py -a gelutanh -q 0  # bf16
python3 op_tests/test_moe_2stage.py -a gelutanh -q 1 --no-flydsl-csv   # per_tensor fp8
python3 op_tests/test_moe_2stage.py -a gelutanh -q 2 --no-flydsl-csv   # per_token  fp8

github-actions · 2026-08-07T06:16:25Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4620 --add-label <label>

Copilot

Pull request overview

This PR adds end-to-end support for the tanh-approximation GELU activation (GeluTanh) in the Composable Kernel (CK) XDL 2-stage MoE path, including C++/pybind enum exposure, CK stage-1 activation mapping, and codegen/test reference wiring.

Changes:

Introduces ActivationType::GeluTanh = 4 and exposes it via pybind.
Maps GeluTanh to CK’s stage-1 act_op=4 (gelu_tanh_and_mul) and wires "gelutanh": 4 into CK 2-stage codegen.
Adds a torch reference implementation (F.gelu(..., approximate="tanh")) and improves CLI activation parsing to be case-insensitive.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
csrc/include/aiter_enum.h	Adds the new `ActivationType::GeluTanh` enum value.
csrc/include/rocm_ops.hpp	Exposes `GeluTanh` through the pybind `ActivationType` enum.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages.cu	Maps AITER activation to CK stage-1 `act_op`, adding `GeluTanh -> 4`.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py	Extends `ACT_OP_MAP`/naming to include `gelutanh: 4`.
csrc/ck_gemm_moe_2stages_codegen/gen_instances.py	Adds `gelutanh` CLI option and an AOT prebuild loop for plain-f8 quant instances.
aiter/utility/dtypes.py	Makes activation parsing case-insensitive for mixed-case enum members like `GeluTanh`.
aiter/ops/quant.py	Adds torch reference mapping for `GeluTanh` using tanh-approx GELU.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

a-sidorova · 2026-08-10T05:40:01Z

 # CK stage1 activation op values. swiglu (OAI swiglu_oai) maps to CK act value 3;
-# value 2 is intentionally skipped (reserved for a future activation variant).
-ACT_OP_MAP = {"gelu": 0, "silu": 1, "swiglu": 3}
+# gelutanh (gelu_tanh_and_mul) maps to CK act value 4; value 2 is intentionally
+# skipped (reserved for a future activation variant).
+ACT_OP_MAP = {"gelu": 0, "silu": 1, "swiglu": 3, "gelutanh": 4}


This can't silently run the wrong activation on the CK 2-stage path: ck_moe_stage1 selects the compiled module by the requested activation, so the loaded module's lookup table only contains that activation's instances (names embed _gelutanh_). A fallback Silu/Gelu kernelName1 won't resolve there, so C++ logs ck kernel not found and falls through to moe_stage1_heuristic_dispatch, which matches on ActOP == act_op and runs the correct GeluTanh epilogue

a-sidorova marked this pull request as ready for review August 7, 2026 11:16

a-sidorova requested review from a team and a lite review from Copilot August 7, 2026 11:16

Copilot started reviewing on behalf of a-sidorova August 7, 2026 11:17 View session

Copilot AI reviewed Aug 7, 2026

View reviewed changes

zufayu requested a review from yzhou103 August 10, 2026 01:16

a-sidorova added 2 commits August 10, 2026 05:45

[MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE

dfca9af

applied copilot's comment for str2ActivationType

a5cf348

a-sidorova force-pushed the asidorov/feature/ck_moe_gelu_tanh branch from 7b28e13 to a5cf348 Compare August 10, 2026 05:46

Navigation Menu

Uh oh!

[MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE - #4620

[MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE#4620
a-sidorova wants to merge 2 commits into
ROCm:mainfrom
a-sidorova:asidorov/feature/ck_moe_gelu_tanh

a-sidorova commented Aug 7, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

a-sidorova Aug 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

a-sidorova commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Uh oh!

github-actions Bot commented Aug 7, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

a-sidorova Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

a-sidorova commented Aug 7, 2026 •

edited

Loading